package MLDA_Process;

import gnu.trove.map.hash.TIntIntHashMap;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import Commons.IOUtil;
import Commons.voseAlias;

public class AliasPAM4ST implements Runnable {

    public int P = 1000;
    public int K = 100;

    public int M;
    public int V;

    public int MH_STEPS = 2;

    public double alpha1 = 0.1;
    public double alpha2 = 0.1;

    public double beta = 0.01;

    public double Vbeta;
    public double Kalpha2;

    public int nwk[][];
    public int nk[];

    public int npk[][];
    public int np[];

    public TIntIntHashMap[] ndp;

    public int niters;

    public List<voseAlias> q = new ArrayList<voseAlias>();

    public int[][] z1_assigns;
    public int[][] z2_assigns;

    public double[] p;

    public List<List<Integer>> docs = new ArrayList<List<Integer>>();
    public HashMap<String, Integer> w2i = new HashMap<String, Integer>();
    public HashMap<Integer, String> i2w = new HashMap<Integer, String>();


    public int saveStep = 1000;
    public String inputPath = "";
    public String outputPath = "";


    public AliasPAM4ST(int P, int K, int niter, int saveStep, double alpha1, double alpha2, double beta, String inputPath, String outputPath) {
        this.P = P;
        this.K = K;
        this.alpha1 = alpha1;
        this.alpha2 = alpha2;
        this.saveStep = saveStep;
        this.beta = beta;
        this.niters = niter;
        this.inputPath = inputPath;
        this.outputPath = outputPath;
    }

    public void loadTxts(String txtPath) {
        BufferedReader reader = IOUtil.getReader(txtPath, "UTF-8");

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();

                String[] tokens = line.trim().split("\\s+");
                for (String token : tokens) {
                    if (!w2i.containsKey(token)) {
                        w2i.put(token, w2i.size());
                        i2w.put(w2i.get(token), token);
                    }
                    doc.add(w2i.get(token));
                }
                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        M = docs.size();
        V = w2i.size();
        return;
    }

    public void setParameters(double alpha1, double alpha2, double beta, int npseudo_doc, int ntopics, int niters) {
        this.P = npseudo_doc;
        this.K = ntopics;
        this.alpha1 = alpha1;
        this.alpha2 = alpha2;
        this.beta = beta;
        this.niters = niters;
    }

    public void init_model() {
        Vbeta = V * beta;
        Kalpha2 = K * alpha2;

        p = new double[K];

        nwk = new int[V][K];
        nk = new int[K];

        npk = new int[P][K];
        np = new int[P];

        ndp = new TIntIntHashMap[M];

        z1_assigns = new int[M][];
        z2_assigns = new int[M][];

        for (int m = 0; m != M; m++) {
            ndp[m] = new TIntIntHashMap();
            int N = docs.get(m).size();
            z1_assigns[m] = new int[N];
            z2_assigns[m] = new int[N];
            for (int n = 0; n != N; n++) {
                int w = docs.get(m).get(n);
                int z1 = (int) Math.floor(Math.random() * P);
                int z2 = (int) Math.floor(Math.random() * K);

                npk[z1][z2]++;
                np[z1]++;

                nwk[w][z2]++;
                nk[z2]++;

                ndp[m].adjustOrPutValue(z1, 1, 1);

                z1_assigns[m][n] = z1;
                z2_assigns[m][n] = z2;
            }
        }

        this.specific_init();
    }

    public void specific_init() {
        System.out.println("Initializing the alias tables ...");
        for (int w = 0; w != V; w++) {
            q.add(new voseAlias());
        }
        for (int w = 0; w < V; ++w) {
            q.get(w).init(P * K);
            generateQtable(w);
        }
    }

    public void generateQtable(int w) {
        q.get(w).wsum = 0.0;
        for (int p = 0; p != P; p++) {
            for (int k = 0; k != K; k++) {
                q.get(w).w[p * K + k] = (npk[p][k] + alpha2) / (np[p] + Kalpha2)
                        * (nwk[w][k] + beta) / (nk[k] + Vbeta);
                q.get(w).wsum += q.get(w).w[p * K + k];
            }
        }
        q.get(w).constructTable();
    }

    public void sample_doc(int m) {
        int N = docs.get(m).size();

        for (int n = 0; n != N; n++) {
            int w = docs.get(m).get(n);

            // remove z_ij from the count variables
            int pseudo_doc = z1_assigns[m][n];
            int topic = z2_assigns[m][n];
            int new_pseudo_doc = -1;
            int new_topic = -1;

            p = new double[ndp[m].keys().length * K];

            nwk[w][topic]--;
            nk[topic]--;

            npk[pseudo_doc][topic]--;
            np[pseudo_doc]--;

            if (ndp[m].get(pseudo_doc) == 1) {
                ndp[m].remove(pseudo_doc);
            } else {
                ndp[m].adjustValue(pseudo_doc, -1);
            }

			/* Travese all non-zero document-topic distribution */
            double psum = 0.0;
            int index = 0;
            for (int pd : ndp[m].keys()) {
                for (int k = 0; k != K; k++) {
                    psum += ndp[m].get(pd) * (npk[pd][k] + alpha2) / (np[pd] + Kalpha2) *
                            (nwk[w][k] + beta) / (nk[k] + Vbeta);
                    p[index] = psum;
                    index++;
                }
            }

            double select_pr = psum / (psum + alpha1 * q.get(w).wsum);

            //MHV to draw new topic
            for (int s = 0; s < MH_STEPS; s++) {
                //1. Flip a coin
                if (Math.random() < select_pr) {
                    double u = Math.random() * psum;
                    for (int i = 0; i != index; i++) {
                        if (u < p[i]) {
                            new_pseudo_doc = ndp[m].keys()[i / K];
                            new_topic = i % K;
                            break;
                        }
                    }
                } else {
                    q.get(w).noSamples++;
                    if (q.get(w).noSamples > (P / (double) MH_STEPS)) {
                        generateQtable(w);
                    }
                    int temp = q.get(w).sample(Math.random(), Math.random());
                    new_pseudo_doc = temp / K;
                    new_topic = temp % K;
                } // if else

                if (topic != new_topic || pseudo_doc != new_pseudo_doc) {
                    //2. Find acceptance probability
                    double temp_old = (npk[pseudo_doc][topic] + alpha2) / (np[pseudo_doc] + Kalpha2)
                            * (nwk[w][topic] + beta) / (nk[topic] + Vbeta);
                    double temp_new = (npk[new_pseudo_doc][new_topic] + alpha2) / (np[new_pseudo_doc] + Kalpha2)
                            * (nwk[w][new_topic] + beta) / (nk[new_topic] + Vbeta);
                    double acceptance = (ndp[m].get(new_pseudo_doc) + alpha1) / (ndp[m].get(pseudo_doc) + alpha1)
                            * temp_new / temp_old
                            * (ndp[m].get(pseudo_doc) * temp_old + alpha1 * q.get(w).w[pseudo_doc * K + topic])
                            / (ndp[m].get(new_pseudo_doc) * temp_new + alpha1 * q.get(w).w[new_pseudo_doc * K + new_topic]);

                    if (Math.random() < acceptance) {
                        pseudo_doc = new_pseudo_doc;
                        topic = new_topic;
                    }
                }
            } // for
            nwk[w][topic]++;
            nk[topic]++;

            npk[pseudo_doc][topic]++;
            np[pseudo_doc]++;

            ndp[m].adjustOrPutValue(pseudo_doc, 1, 1);

            z1_assigns[m][n] = pseudo_doc;
            z2_assigns[m][n] = topic;
        }
    }

    public void estimate() {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("Alias PAM4ST Interation " + iter + " ...");
            if (iter % this.saveStep == 0 && iter != 0 && iter != this.niters - 1) {
                this.storeResult(iter);
            }
            long begin = System.currentTimeMillis();
            for (int m = 0; m != M; m++) {
                ndp[m].compact();
                this.sample_doc(m);
            }
            System.out.println((System.currentTimeMillis() - begin) / 1000.0);
        }
    }

    public double[][] computePhi() {
        double[][] phi = new double[K][V];
        for (int k = 0; k != K; k++) {
            for (int v = 0; v != V; v++) {
                phi[k][v] = (nwk[v][k] + beta) / (nk[k] + Vbeta);
            }
        }
        return phi;
    }

    public int[][] ndk;
    public int[] nd;
    public double alpha = 0.1;

    public void convert_zassigns_to_arrays_theta() {
        ndk = new int[M][K];
        nd = new int[M];

        for (int m = 0; m != M; m++) {
            for (int n = 0; n != docs.get(m).size(); n++) {
                ndk[m][z2_assigns[m][n]]++;
                nd[m]++;
            }
        }
    }

    public double[][] computeTheta() {
        this.convert_zassigns_to_arrays_theta();
        double[][] theta = new double[M][K];
        for (int m = 0; m != M; m++) {
            for (int k = 0; k != K; k++) {
                theta[m][k] = (ndk[m][k] + alpha) / (nd[m] + K * alpha);
            }
        }
        return theta;
    }

    public ArrayList<List<Entry<String, Double>>> sortedTopicWords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : w2i.keySet())
                term2weight.put(term, phi[k][w2i.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });

            res.add(pairs);
        }

        return res;
    }

    public void printTopics(String path, int top_n) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(path);
        double[][] phi = computePhi();
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K);
        for (int k = 0; k != K; k++) {
            writer.write("Topic " + k + ":\n");
            for (int i = 0; i != top_n; i++) {
                writer.write(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue() + "\n");
            }
        }
    }

    public void savePhi(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] phi = computePhi();
        try {
            for (int k = 0; k != K; k++) {
                for (int v = 0; v != V; v++) {
                    writer.append(phi[k][v] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveTheta(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] theta = computeTheta();
        try {
            for (int m = 0; m != M; m++) {
                for (int k = 0; k != K; k++) {
                    writer.append(theta[m][k] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveWordmap(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (String word : w2i.keySet())
                writer.append(word + "\t" + w2i.get(word) + "\n");

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void printModel() {
        System.out.println("P :" + this.P +
                "\tK :" + this.K +
                "\tniters :" + this.niters +
                "\tsaveStep :" + this.saveStep +
                "\talpha1 :" + this.alpha1 +
                "\talpha2 :" + this.alpha2 +
                "\tbeta :" + this.beta +
                "\tinputPath :" + this.inputPath +
                "\toutputPath :" + this.outputPath);
    }

    public void storeResult(int times) {
        String appendString = "final";
        if (times != 0) {
            appendString = times + "";
        }
        try {
            this.printTopics(outputPath + "/model-" + appendString + ".twords", 20);
            this.saveWordmap(outputPath + "/wordmap.txt");
            this.savePhi(outputPath + "/model-" + appendString + ".phi");
            //this.saveZ2Assigns(outputPath+"/model-"+appendString+".tassign");
            this.saveTheta(outputPath + "/model-" + appendString + ".theta");
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
    }

    public void run() {
        printModel();
        this.loadTxts(inputPath);//
        this.init_model();
        this.estimate();
        this.storeResult(0);

    }

    public static void AliasPAMForShortText(int P, int K, int iter, int saveStep, double alpha1, double alpha2, double beta, int threadNum, String path) {
        ExecutorService service = Executors.newCachedThreadPool();
        File trainFile = new File(path);
        String parent_path = trainFile.getParentFile().getAbsolutePath() + "/AliasPAM4ST_" + P + "_" + K + "_" + alpha1 + "_" + alpha2 + "_" + beta + "/";
        for (int i = 0; i < threadNum; i++) {
            (new File(parent_path + i + "/")).mkdirs();
            service.submit(new AliasPAM4ST(P, K, iter, saveStep, alpha1, alpha2, beta, path, parent_path + i + "/"));

        }
        service.shutdown();
    }

    // public static void main(String args[]) {
    // 	AliasPAM4ST tm = new AliasPAM4ST();
    // 	tm.loadTxts("/Users/zuoyuan/Desktop/Experiment/20130601-origin");//
    // 	tm.setParameters(0.001, 0.5, 0.01, 200, 20, 500);
    // 	tm.init_model();
    // 	tm.estimate();
    // 	tm.printTopics(10);
    // }

}
